import talib
import numpy as np


class IndicatorControl():
    """
    指标中心
    """

    def latestMaValue(self,
                      closeArr,
                      time,
                      maType=talib.MA_Type.SMA,
                      size=None):
        """
        最新的移动平均线价格
        """
        if size == None or size == 0:
            return talib.MA(np.array(closeArr), time, maType)[-1]
        return talib.MA(np.array(closeArr), time, maType)[-size:]

    def maValueForTimeArray(self, closeArr, timeArray, maType=0, size=None):
        """
        多个移动平均线价格
        """
        value = []
        maLineType = talib.MA_Type.SMA
        if maType == 1:
            maLineType = talib.MA_Type.EMA

        for time in timeArray:
            value.append(
                self.latestMaValue(closeArr,
                                   time,
                                   maType=maLineType,
                                   size=size))
        return value

    def rsiValueForMoreTimeFrame(self, closeArrList):
        """
        RSI指标数据
        """
        valueArr = []
        for item in closeArrList:
            valueArr.append(talib.RSI(np.array(item))[-1])
        return valueArr

    def rocValueForMoreTimeFrame(self, closeArrList):
        """
        ROC指标数据
        """
        valueArr = []
        for item in closeArrList:
            valueArr.append(talib.ROC(np.array(item), timeperiod=14)[-1])
        return valueArr

    def macdValueForMoreTimeFrame(self, closeArrList):
        """
        MACD指标数据
        """
        valueArr = []
        for item in closeArrList:
            macd, signal, hist = talib.MACD(np.array(item),
                                            fastperiod=12,
                                            slowperiod=26,
                                            signalperiod=9)
            valueArr.append(macd[-1])
        return valueArr

    def ultoscValueForMoreTimeFrame(self, priceArrList):
        """
        ULTOSC指标数据
        """
        valueArr = []
        for item in priceArrList:
            ultosc = talib.ULTOSC(np.array(item[1]), np.array(item[2]),
                                  np.array(item[0]))
            valueArr.append(ultosc[-1])
        return valueArr

    def willrValueForMoreTimeFrame(self, priceArrList):
        """
        WILLR指标数据
        """
        valueArr = []
        for item in priceArrList:
            willr = talib.WILLR(np.array(item[1]),
                                np.array(item[2]),
                                np.array(item[0]),
                                timeperiod=24)
            valueArr.append(willr[-1])
        return valueArr

    def stochValueForMoreTimeFrame(self, priceArrList):
        """
        STOCH指标数据
        """
        valueArr = []
        for item in priceArrList:
            STOCH_slowk, STOCH_slowd = talib.STOCH(np.array(item[1]),
                                                   np.array(item[2]),
                                                   np.array(item[0]),
                                                   fastk_period=9,
                                                   slowk_period=6,
                                                   slowk_matype=0,
                                                   slowd_period=3,
                                                   slowd_matype=0)
            valueArr.append(STOCH_slowd[-1])
        return valueArr

    def adxValueForMoreTimeFrame(self, priceArrList):
        """
        ADX指标数据
        """
        valueArr = []
        for item in priceArrList:
            adx = talib.ADX(np.array(item[1]),
                            np.array(item[2]),
                            np.array(item[0]),
                            timeperiod=14)
            plus_di = talib.PLUS_DI(np.array(item[1]),
                                    np.array(item[2]),
                                    np.array(item[0]),
                                    timeperiod=14)
            minus_di = talib.MINUS_DI(np.array(item[1]),
                                      np.array(item[2]),
                                      np.array(item[0]),
                                      timeperiod=14)

            if plus_di[-1] > minus_di[-1] and plus_di[-1] > adx[-1]:
                valueArr.append(True)
            elif minus_di[-1] > plus_di[-1] and minus_di[-1] > adx[-1]:
                valueArr.append(False)
            else:
                valueArr.append(None)
        return valueArr

    def cciValueForMoreTimeFrame(self, priceArrList):
        """
        CCI指标数据
        """
        valueArr = []
        for item in priceArrList:
            cci = talib.CCI(np.array(item[1]),
                            np.array(item[2]),
                            np.array(item[0]),
                            timeperiod=20)
            valueArr.append(cci[-1])
        return valueArr

    def atrValueForMoreTimeFrame(self, priceArrList):
        """
        ATR指标数据
        """
        valueArr = []
        for item in priceArrList:
            atr = talib.ATR(np.array(item[1]),
                            np.array(item[2]),
                            np.array(item[0]),
                            timeperiod=14)
            valueArr.append(atr[-1])
        return valueArr

    def stochRsiValueForMoreTimeFrame(self, closeArrList):
        """
        STOCHRSI指标数据
        """
        valueArr = []
        for item in closeArrList:
            STOCH_slowk, STOCH_slowd = talib.STOCHRSI(np.array(item),
                                                      timeperiod=14,
                                                      fastk_period=3,
                                                      fastd_period=3,
                                                      fastd_matype=0)
            valueArr.append(STOCH_slowd[-1])
        return valueArr